import torch
from torch import nn                         #完成神经网络的构建包
from torch.nn import functional as F         #包含常用的函数包
from torch import optim                      #优化工具包
import torchvision                           #视觉工具包
import  matplotlib.pyplot as plt            #画图工具
from python_utils import plot_curve,plot_image,one_hot
import os                                    #允许副本存在
os.environ['KMP_DUPLICATE_LIB_OK']='True'

batch_size=512

#step1 load dataset   加载数据包
# mnist_data'：加载mnist数据集，路径
# train=True：选择训练集还是测试
# download=True：如果当前文件没有mnist文件就会自动从网上去下载
# torchvision.transforms.ToTensor():下载好的数据一般是numpy格式，转换成Tensor
# torchvision.transforms.Normalisze((0.1307,), (0.3081,))：正则化过程，为了让数据更好的在0的附近均匀的分布
# 均值是0.1307，标准差是0.3081，这些系数都是数据集提供方计算好的数据
train_loader=torch.utils.data.DataLoader(
    torchvision.datasets.MNIST("minist_data",train=True,download=True,transform=torchvision.transforms.Compose(
        [torchvision.transforms.ToTensor(),torchvision.transforms.Normalize((0.1307,),(0.3081,))
         ])),
    batch_size=batch_size,shuffle=True)
    # batch_size=batch_size：表示一次加载多少张图片
    # shuffle = True 加载的时候做一个随机的打散

test_loader=torch.utils.data.DataLoader(
    torchvision.datasets.MNIST("minist_data",train=True,download=False,transform=torchvision.transforms.Compose(
        [torchvision.transforms.ToTensor(),torchvision.transforms.Normalize((0.1307,),(0.3081,))
         ])),
    batch_size=batch_size,shuffle=False)

# x=>torch.Size([512, 1, 28, 28])代表512张1通道28*28的灰度图
# y=>torch.Size([512])代表512张图片类别
x,y=next(iter(train_loader))
print(x.shape,y.shape)
plot_image(x,y,"image")
print(x)
print(y)

#step2 构建神经网络结构 ，一个三层的线性全连接网络
class Net(nn.Module):  # 继承 torch 的 Module（固定）

    def __init__(self):
        super(Net,self).__init__()  # 继承 __init__ 功能（固定）
        # 定义每层用什么样的形式, 28*28=784是输入， 10是10种分类，这两个是固定的。 256和64由经验决定
        #xw+b
        self.fc1=nn.Linear(28*28,256)# 定义fc1层，28*28输入，256线性输出
        self.fc2=nn.Linear(256,64)# 定义fc2层，256输入，64线性输出
        self.fc3=nn.Linear(64,10)# 定义fc3层，64输入，10线性输出

    # x是输入信息就是data，同时也是 Module 中的 forward 功能，
    # 定义神经网络前向传递的过程
    # 把__init__中的层信息一个一个的组合起来
    # 正向传播输入值, 神经网络分析出输出值
    def forward(self, x):
        #x:[b,1,28,28]
        #h1=relu(xw1+b1)
        x=F.relu(self.fc1(x))  # 定义激励函数(隐藏层的线性值)
        #h2=relu(h1w2+b2)
        x=F.relu(self.fc2(x))
        #h3=h2w3+b3
        x=(self.fc3(x))   # 分类问题一般是softmax + mean square error均方差 ，简单起见直接使用softmax
        return x

net=Net()
# print(net)
# 神经网络结构如下
# Net(
#   (fc1): Linear(in_features=784, out_features=256, bias=True)
#   (fc2): Linear(in_features=256, out_features=64, bias=True)
#   (fc3): Linear(in_features=64, out_features=10, bias=True)
# )

optimizer=optim.SGD(net.parameters(),lr=0.01,momentum=0.9)
# 1.optimizer神经网络优化器
# 2.优化方法：sgd/momentum/Nesterov/adagrad/adadelta
# 3.SGD 是最普通的优化器, 也可以说没有加速效果, 而 Momentum 是 SGD 的改良版, 它加入了动量原则.
# 后面的 RMSprop 又是 Momentum 的升级版. 而 Adam 又是 RMSprop 的升级版
# 4.Stochastic Gradient Descent (SGD)随机梯度下降
# 5.Momentum 传统的参数 W 的更新是把原始的 W 累加上一个负的学习率(learning rate) 乘以校正值 (dx).
# 6.这里的采取momentum 动量加速,在SGD函数里指定momentum的值，lr学习率
# 7.net.parameters(): 会帮我们拿到权值 [w1, b1, w2, b2, w3, b3]



train_loss=[]  # 定义存贮损失函数loss值列表

#  step3，开始训练，将数据集遍历三遍
for epoch in range(3):

    for batch_idx,(x,y) in enumerate(train_loader):  # 对整个训练集迭代一次
        # 这里是对一个batch迭代一次，一次batch 512张图片
        # 输入 torch.Size([512, 1, 28, 28]) 输出 torch.Size([512])
        # x: [b, 1, 28, 28], y: [512]
        # 需要把输入tensor[b, 1, 28, 28] -> 打平为下面[b, feature]
        # [b, 1, 28, 28] => [b, 784]
        x=x.view(x.size(0),28*28)  # x.size(0)表示batch 512    (512, 1, 28, 28) -> (512, 784)
        # =>512行784列矩阵
        out =net(x)  # 经过了class Net(nn.Module)，输入512张784像素的图片

        y_onehot=one_hot(y)  # one_hot编码类别在[0,9]
        #loss=mse(out,y_onehot)  mse 是 均方差 mean square error
        loss= F.mse_loss(out,y_onehot)  #  计算out与y_onehot之间的均方差，得到loss

        optimizer.zero_grad()  # 先对梯度进行清零
        loss.backward()  # 梯度计算过程，计算梯度
        #w'=w-lr*grad
        optimizer.step()  #更新权值

        train_loss.append(loss.item())  # 存储损失值

        if batch_idx %10==0:  # 每隔10个batch打印一下
            print(epoch,batch_idx,loss.item())   # 第几个大循环（一共3个）， 第多少批次eg:10 20 30 , loss显示

# step4：画出训练过程的loss曲线
plot_curve(train_loss)
# get optimal [w1,b1,w2,b2,w3,b3]

# step5，检验测试集的准确率
total_correct=0
for x,y in test_loader:
    x=x.view(x.size(0),28*28)
    out=net(x)
    # out: [b, 10] => pred: [b]
    pred=out.argmax(dim=1)  #取得[b, 10]的10个值的最大值所在位置的索引
    correct=pred.eq(y).sum().float().item()  #满足预测值与标签相等的张量值化为float相加
    total_correct+=correct

total_num=len(test_loader.dataset)
acc=total_correct / total_num  # test acc: 0.7799
print("test.acc:",acc)  #输出整体预测的准确度

# 取一个batch,查看预测结果
x,y=next(iter(test_loader))
out=net(x.view(x.size(0),28*28))
pred=out.argmax(dim=1)
plot_image(x,pred,"test")


